#! usr/bin/env python3
# -*- coding: utf-8 -*-

from gmpy2 import mpz, is_prime, random_state, mpz_urandomb, mpz_random, powmod, bit_length, mul, invert, t_mod, \
    next_prime, f_div
import random
import sys
import hashlib

import AES

rand = random_state(random.randrange(sys.maxsize))


class KGC(object):
    def __init__(self, bits):
        self.p_q_gen(bits)
        self.createpk()
        self.createsk()
        self.createP()

    def p_q_gen(self, bits):
        self.p = mpz(2) ** (bits - 1) + mpz_urandomb(rand, bits - 1)
        self.p = next_prime(self.p)
        self.q = mpz(2) ** (bits - 1) + mpz_urandomb(rand, bits - 1)
        self.q = next_prime(self.q)

    def createpk(self):
        self.pk = mpz(65537)

    def createsk(self):
        self.sk = invert(self.pk, (self.p - 1) * (self.q - 1))

    def createP(self):
        self.P = self.p * self.q


class Params(object):
    def __init__(self, pa, pb, bits):
        self.pa = pa
        self.pb = pb
        self.bits = bits


# user key class
class User(object):

    def __init__(self, uid, pk, pa, bits):
        self.pk = pk
        self.id = str(H1_hash(uid, pk, pa, bits))

    # signcryption function
    def signcrypt(self, Reciever, params, m, sk):
        w = mpz_urandomb(rand, params.bits - 1)
        key = mpz_urandomb(rand, 128)
        localkey = str(key).encode("utf-8")
        # print(localkey)
        aeskey = hash_func(localkey)
        aeskey = bytes.fromhex(aeskey[0:32])
        CText = AES.FileAES(aeskey)
        cipher_text = CText.encrypt(m)
        h1 = self.H1(self.id, Reciever.id, w, params.bits)
        t = self.T(w, Reciever.pk, params.pb)

        c = self.C(key, w, params.pb)
        h2 = self.H2(self.id, t, key, c, params.bits)
        s = self.S(h2, sk, params.pa)
        Signcryption_text = (h1, t, c, s,cipher_text)
        return Signcryption_text

    def unsigncrypt(self, Sender, params, Signcryption_text, sk):
        h1 = Signcryption_text[0]
        t = Signcryption_text[1]
        c = Signcryption_text[2]
        s = Signcryption_text[3]
        ct=Signcryption_text[4]
        w = powmod(t, sk, params.pb)

        key = t_mod(c * invert(w, params.pb), params.pb)
        m=self.deaes(key,ct)
        h2 = self.H2(Sender.id, t, key, c, params.bits)
        h2_be = powmod(s, Sender.pk, params.pa)
        if (h1 == self.H1(Sender.id, self.id, w, params.bits) and h2 == h2_be):
            return m
        else:
            return False

    def deaes(self,key,ct):
        localkey = str(key).encode("utf-8")
        # print(localkey)
        aeskey = hash_func(localkey)
        aeskey = bytes.fromhex(aeskey[0:32])
        # print("aeskey:",aeskey)
        aesde = AES.FileAES(aeskey)
        m = aesde.decrypt(ct)
        return m


    def H1(self, uid1, uid2, w, bits):
        return H1_hash(uid1, uid2, w, bits)

    def T(self, w, pkb, Pb):
        return powmod(w, pkb, Pb)

    def C(self, m, w, p):
        return t_mod(m * w, p)

    def H2(self, uid1, T, m, C, bits):
        return H2_hash(uid1, T, m, C, bits)

    def S(self, h2, ska, pa):
        return powmod(h2, ska, pa)


# define the hashfunction that return 2048 bits hashvalue
def hash_func(x):
    temp1 = hashlib.sha512()
    temp2 = hashlib.sha512()
    temp3 = hashlib.sha512()
    temp4 = hashlib.sha512()
    temp1.update(x + b"1")
    temp2.update(x + b"2")
    temp3.update(x + b"3")
    temp4.update(x + b"4")
    ret_str = temp1.hexdigest() + temp2.hexdigest() + temp3.hexdigest() + temp4.hexdigest()
    return ret_str


# define the H1 hashfunction
def H1_hash(uid, x, y, bits):
    temp = hash_func(uid.encode('utf-8') + str(x).encode('utf-8') + str(y).encode('utf-8'))
    H1_hash_value = mpz(temp[: bits - 2], 16)
    return H1_hash_value


# define the H2 hashfunction
def H2_hash(uid, x, y, z, bits):
    temp = hash_func(uid.encode('utf-8') + str(x).encode('utf-8') + str(y).encode('utf-8') + str(z).encode('utf-8'))
    bits = bits / 2 - 2
    H2_hash_value = mpz(temp[: int(bits)], 16)
    return H2_hash_value


# define the H3 hashfunction
def H3_hash(x, bits):
    temp = hash_func(str(x).encode('utf-8'))
    H3_hash_value = mpz(temp[: bits - 2], 16)
    return H3_hash_value


if __name__ == "__main__":
    alicedone = KGC(512)
    bobdone = KGC(512)
    params = Params(alicedone.P, bobdone.P, 512)
    # print(alicedone.P,bobdone.P)
    alice = User("Alice", alicedone.pk, alicedone.P, params.bits)
    bob = User("Bob", bobdone.pk, bobdone.P, params.bits)
    m = "hello world!"
    Signcryption_text = alice.signcrypt(bob, params, m, alicedone.sk)
    # print(Signcryption_text)
    m = bob.unsigncrypt(alice, params, Signcryption_text, bobdone.sk)
    print(m)